import openai
from .base import BaseProvider

class OpenAIProvider(BaseProvider):
    def __init__(self, base_url=None, api_key=None, timeout=60):
        super().__init__(base_url, api_key, timeout)
        # 可在此处初始化openai配置
        if api_key:
            openai.api_key = api_key
        if base_url:
            openai.api_base = base_url

    def chat(self, model: str, prompt: str, **kwargs) -> str:
        """调用大语言模型生成回复"""
        openai.api_key = self.api_key
        if self.base_url:
            openai.api_base = self.base_url
        resp = openai.ChatCompletion.create(
            model=model,
            messages=[{"role": "user", "content": prompt}],
            timeout=self.timeout,
            **kwargs
        )
        return resp.choices[0].message["content"]

    def embed(self, model: str, text: str, **kwargs) -> list[float]:
        """调用向量嵌入模型"""
        openai.api_key = self.api_key
        if self.base_url:
            openai.api_base = self.base_url
        resp = openai.Embedding.create(
            model=model,
            input=text,
            timeout=self.timeout,
            **kwargs
        )
        return resp.data[0].embedding
    
    def chat_stream(self, model: str, prompt: str, **kwargs):
        """调用大语言模型流式生成回复，返回一个迭代器，每次yield一小段输出"""
        import requests
        url = f"{self.base_url}/chat/completions"
        headers = {"Authorization": f"Bearer {self.api_key}"}
        data = {
            "model": model,
            "messages": [{"role": "user", "content": prompt}],
            "stream": True,
            **kwargs
        }
        try:
            with requests.post(url, headers=headers, json=data, stream=True, timeout=self.timeout) as r:
                r.raise_for_status()
                for line in r.iter_lines():
                    if line:
                        # OpenAI流式返回一般是json格式，建议解析后yield内容
                        try:
                            import json
                            chunk = json.loads(line.decode("utf-8"))
                            yield chunk.get("choices", [{}])[0].get("delta", {}).get("content", "")
                        except Exception:
                            yield line.decode("utf-8")
        except Exception as e:
            yield f"[ERROR] {str(e)}"
